from langchain.chat_models import init_chat_model
from langchain.prompts import (
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    ChatPromptTemplate,
)


llm = init_chat_model(
    model_provider="openai",
    model="qwen/qwq-32b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key="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",
)


sys_tmpl = """
根据用户输入内容，分析最适合的数据可视化类型。请从类型中选择最匹配的一项：

可选类型：
1. 折线图 - 适合展示数据随时间或有序类别的趋势变化
2. 柱状图 - 适合比较不同类别之间的数值大小
3. 饼图 - 适合展示整体中各部分的比例关系
4. 纯文本 - 当数据不适合可视化或只需简单展示时使用

请严格按照以下JSON格式返回结果，只需给出JSON内容，不需要其他说明：
{{
    "type": "选择的具体类型名称",
    "desc": "类型描述",
    "reason": "简要说明选择该类型的理由"
}}

注意事项：
- 如果用户明确要求某种图表类型，优先尊重用户选择
- 当数据包含时间序列时优先考虑折线图
- 当需要比较不同类别数值时优先考虑柱状图
- 当需要显示占比或组成关系时使用饼图
- 当数据量很小或用户明确要求文字展示时使用纯文本
"""

print("sys_tmpl", sys_tmpl)

# 创建消息模板
sys_msg = SystemMessagePromptTemplate.from_template(sys_tmpl)
user_msg = HumanMessagePromptTemplate.from_template("""
已成功根据文档ID：'1600bda5-47aa-4862-b578-e50f48fbf413'生成饼图配置信息，饼图配置ID：1600bda5-47aa-4862-b578-e50f48fbf413
""")

# 组合聊天提示
chat_prompt = ChatPromptTemplate.from_messages([sys_msg, user_msg])

# 格式化提示并调用模型
formatted_prompt = chat_prompt.format()
response = llm.invoke(formatted_prompt)

print(response.content)
# {
#     "type": "饼图",
#     "desc": "适合展示整体中各部分的比例关系",
#     "reason": "用户已明确生成饼图配置，说明需求符合展示占比或组成关系的场景"
# }
